# -*- coding: utf-8 -*-
"""
Created on Fri Jan  6 16:57:59 2023

@author: 18735
"""
import os
import cv2
import matplotlib.pyplot as plt
from skimage import morphology
import SimpleITK as sitk
import numpy as np
import time


def thicker(image, thickness=2):
    thick_image = image.copy()
    for i in range(image.shape[0]):
        for j in range(image.shape[1]):
            if image[i, j]:
                label_value = image[i, j]
                for off in range(-thickness, thickness + 1):
                    try:
                        thick_image[i, j + off] = label_value
                    except:
                        pass
                for off in range(-thickness, thickness + 1):
                    try:
                        thick_image[i + off, j] = label_value
                    except:
                        pass
    return thick_image


def get_exist_range(image_array):
    start, end = 0, 0
    for i in range(image_array.shape[0]):
        if image_array[i].any():
            start = i
            break
    for i in range(image_array.shape[0] - 1, -1, -1):
        if image_array[i].any():
            end = i + 1
            break
    return start, end


def get_center(image_array: np.ndarray, ratio_down=1/3, ratio_up=2/3) -> tuple:
    """
    calculate the rotation angle and center from the XY slices

    :param image_array: 3D array seen as ZXY
    :param ratio_down: lower bound along the Z axis
    :param ratio_up:
    :return: point(XYZ), angle(along z axis)
    """
    start, end = get_exist_range(image_array)
    length = end - start
    temp = start
    start = temp + ratio_down * length
    end = temp + ratio_up * length

    center = [0., 0.]
    cnt = 0
    for i in range(round(start), round(end)):
        slice = image_array[i]
        slice = slice.astype(np.uint8)
        slice = slice[:, :, np.newaxis]
        contours, _ = cv2.findContours(slice, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
        rect = cv2.minAreaRect(contours[0])
        center[0] += rect[0][0]
        center[1] += rect[0][1]
        cnt += 1

    center[0] /= cnt
    center[1] /= cnt
    return center[0], center[1], (start + end) / 2


def get_skeleton(label_array, max_num_iter=15, thickness=2):
    """
    calculate skeleton by traversing slices along Z-axis.

    :param label_array: ZYX array
    :param thickness: thickness of the skeleton
    :return:
    """
    label_array = label_array.astype(np.uint8)
    skele_array = np.zeros(label_array.shape)
    for z0 in range(label_array.shape[0]):
        slice = label_array[z0]
        if slice.any():
            skeleton = morphology.thin(slice, max_num_iter=max_num_iter)
            skele_array[z0] = thicker(skeleton, thickness)

    return skele_array


def fill_hole_closing(array, radius):
    """
    fill the holes by disk closing along Z, Y and X axis

    :param array: skeleton to be closed
    :param radius: radius of footprint(closing kernal)
    :return:
    """
    ret_array = array.copy()
    k = morphology.disk(radius)
    t = time.process_time()
    t1 = t
    for z0 in range(ret_array.shape[0]):
        ret_array[z0] = morphology.closing(ret_array[z0], k)
    tt = time.process_time()
    print('----finish z axis, cost', tt - t, 'seconds----')
    t = tt
    for y0 in range(ret_array.shape[1]):
        ret_array[:, y0, :] = morphology.closing(ret_array[:, y0, :], k)
    tt = time.process_time()
    print('----finish y axis, cost', tt - t, 'seconds----')
    t = tt
    for x0 in range(ret_array.shape[2]):
        ret_array[:, :, x0] = morphology.closing(ret_array[:, :, x0], k)
    tt = time.process_time()
    print('----finish x axis, cost', tt - t, 'seconds----')
    print('=======total cost:', tt - t1, 'seconds=======')
    return ret_array


def extract_skeleton(image_path, out_path, disk_radius=15, blur_ksize=9, thin_iter_num=15, thickness=4, input_anti=False):
    """
    extract skeleton from input image. If the input is original label, input_anti should be False.

    :param image_path:
    :param out_path:
    :param disk_radius: size of 2D closing footprint. 值越小，运行越快，但不一定可以填上足够大的洞
    :param blur_ksize: kernel size of blurring window. Should always be smaller than disk_radius.
    :param thin_iter_num: max iter num of thinning.
    :param thickness:
    :param input_anti:
    :return:
    """

    image = sitk.ReadImage(image_path)
    array = sitk.GetArrayFromImage(image)
    array = array.astype(np.uint8)

    if not input_anti:
        k = morphology.disk(disk_radius)
        t = time.process_time()
        for z0 in range(array.shape[0]):
            array[z0, :, :] = morphology.closing(array[z0, :, :], k)
        for y0 in range(array.shape[1]):
            array[:, y0, :] = morphology.closing(array[:, y0, :], k)
        for x0 in range(array.shape[2]):
            array[:, :, x0] = morphology.closing(array[:, :, x0], k)
        tt = time.process_time()
        print('finished closing, cost', tt - t, 'seconds')
        # sitk.WriteImage(sitk.GetImageFromArray(array), 'test_closing_before_anti.nii.gz')

        for z0 in range(array.shape[0]):
            array[z0, :, :] = cv2.medianBlur(array[z0, :, :], blur_ksize)
        for y0 in range(array.shape[1]):
            array[:, y0, :] = cv2.medianBlur(array[:, y0, :], blur_ksize)
        for x0 in range(array.shape[2]):
            array[:, :, x0] = cv2.medianBlur(array[:, :, x0], blur_ksize)
        # sitk.WriteImage(sitk.GetImageFromArray(array), 'test_anti.nii.gz')

    array = get_skeleton(array, max_num_iter=thin_iter_num, thickness=thickness)

    k = morphology.disk(10)
    t = time.process_time()
    for y0 in range(array.shape[1]):
        array[:, y0, :] = morphology.closing(array[:, y0, :], k)
    for x0 in range(array.shape[2]):
        array[:, :, x0] = morphology.closing(array[:, :, x0], k)
    tt = time.process_time()
    print('finished closing after skeletonize, cost', tt - t, 'seconds')

    skele_image = sitk.GetImageFromArray(array)
    skele_image.CopyInformation(image)
    sitk.WriteImage(skele_image, out_path)


if __name__ == '__main__':
    image_path = r'C:\study\FFR\myProject\mytest\mydata_apply\input\case_00217\segmentation.nii.gz'
    out_path = r'test_extract_00217.nii.gz'

    extract_skeleton(image_path, out_path)

